import numpy as np
from teleop.open_television.television import TeleVision
from teleop.open_television.constants import *
from teleop.utils.mat_tool import mat_update, fast_mat_inv, rotation_matrix_to_euler_angles, print_pose_info
from teleop.utils.logger import get_logger

# 获取日志记录器
logger = get_logger('TeleVisionWrapper')

"""
(basis) OpenXR Convention : y up, z back, x right. 
(basis) Robot  Convention : z up, y left, x front.  
p.s. Vuer's all raw data follows OpenXR Convention, WORLD coordinate.

under (basis) Robot Convention, wrist's initial pose convention:

    # (Left Wrist) XR/AppleVisionPro Convention:
        - the x-axis pointing from wrist toward middle.
        - the y-axis pointing from index toward pinky.
        - the z-axis pointing from palm toward back of the hand.

    # (Right Wrist) XR/AppleVisionPro Convention:
        - the x-axis pointing from wrist toward middle.
        - the y-axis pointing from pinky toward index.
        - the z-axis pointing from palm toward back of the hand.
  
    # (Left Wrist URDF) Unitree Convention:
        - the x-axis pointing from wrist toward middle.
        - the y-axis pointing from palm toward back of the hand.
        - the z-axis pointing from pinky toward index.

    # (Right Wrist URDF) Unitree Convention:
        - the x-axis pointing from wrist toward middle.
        - the y-axis pointing from back of the hand toward palm. 
        - the z-axis pointing from pinky toward index.

under (basis) Robot Convention, hand's initial pose convention:

    # (Left Hand) XR/AppleVisionPro Convention:
        - the x-axis pointing from wrist toward middle.
        - the y-axis pointing from index toward pinky.
        - the z-axis pointing from palm toward back of the hand.

    # (Right Hand) XR/AppleVisionPro Convention:
        - the x-axis pointing from wrist toward middle.
        - the y-axis pointing from pinky toward index.
        - the z-axis pointing from palm toward back of the hand.

    # (Left Hand URDF) Unitree Convention:   
        - The x-axis pointing from palm toward back of the hand. 
        - The y-axis pointing from middle toward wrist.
        - The z-axis pointing from pinky toward index.

    # (Right Hand URDF) Unitree Convention: 
        - The x-axis pointing from palm toward back of the hand. 
        - The y-axis pointing from middle toward wrist.
        - The z-axis pointing from index toward pinky. 

    p.s. From website: https://registry.khronos.org/OpenXR/specs/1.1/man/html/openxr.html.
         You can find **(Left/Right Wrist) XR/AppleVisionPro Convention** related information like this below:
           "The wrist joint is located at the pivot point of the wrist, which is location invariant when twisting the hand without moving the forearm. 
            The backward (+Z) direction is parallel to the line from wrist joint to middle finger metacarpal joint, and points away from the finger tips. 
            The up (+Y) direction points out towards back of the hand and perpendicular to the skin at wrist. 
            The X direction is perpendicular to the Y and Z directions and follows the right hand rule."
         Note: The above context is of course under **(basis) OpenXR Convention**.

    p.s. **(Wrist/Hand URDF) Unitree Convention** information come from URDF files.
"""

class TeleVisionWrapper:
    def __init__(self, binocular, img_shape, img_shm_name):
        self.tv = TeleVision(binocular, img_shape, img_shm_name)
       

    def get_data(self):

        # --------------------------------wrist-------------------------------------

        # VR坐标系是xyz:右上后
        head_vuer_mat, head_flag = mat_update(const_head_vuer_mat, self.tv.head_matrix())
        left_wrist_vuer_mat, left_wrist_flag  = mat_update(const_left_wrist_vuer_mat, self.tv.left_hand)
        right_wrist_vuer_mat, right_wrist_flag = mat_update(const_right_wrist_vuer_mat, self.tv.right_hand)
        # logger.info(f"------000 header matrix:\n{print_pose_info(self.tv.head_matrix())}")
        # logger.info(f"------000 left wrist matrix:\n{print_pose_info(self.tv.left_hand)}")
        # logger.info(f"------000 right wrist matrix:\n{print_pose_info(self.tv.right_hand)}")
        logger.info(f"------111 head_vuer_mat\n{print_pose_info(head_vuer_mat)}")
        logger.info(f"------111 left_wrist_vuer_mat\n{print_pose_info(left_wrist_vuer_mat)}")
        logger.info(f"------111 right_wrist_vuer_mat\n{print_pose_info(right_wrist_vuer_mat)}")
        
        # 从共享内存中获取控制器状态
        controller_data = {
            # 基本按钮状态
            'right_squeeze': False,
            'right_squeeze_value': 0.0,
            'left_squeeze': False,
            'left_squeeze_value': 0.0,
            # 添加其他按钮状态
            'right_trigger': False,
            'right_trigger_value': 0.0,
            'right_touchpad': False,
            'right_touchpad_x': 0.0,
            'right_touchpad_y': 0.0,
            'right_thumbstick': False,
            'right_thumbstick_x': 0.0,
            'right_thumbstick_y': 0.0,
            'right_a_button': False,
            'right_b_button': False,
            'left_trigger': False,
            'left_trigger_value': 0.0,
            'left_touchpad': False,
            'left_touchpad_x': 0.0,
            'left_touchpad_y': 0.0,
            'left_thumbstick': False,
            'left_thumbstick_x': 0.0,
            'left_thumbstick_y': 0.0,
            'left_x_button': False,
            'left_y_button': False
        }
        
        try:
            # 使用TeleVision类中的方法获取控制器状态
            if hasattr(self.tv, 'right_controller'):
                right_controller_data = self.tv.right_controller()
                # 获取完整的右手柄状态
                if len(right_controller_data) >= 12:  # 确保数组足够长
                    # 基本按钮
                    controller_data['right_squeeze'] = bool(right_controller_data[0] > 0.5)
                    controller_data['right_squeeze_value'] = right_controller_data[1]
                    # 其他按钮
                    controller_data['right_trigger'] = bool(right_controller_data[2] > 0.5)
                    controller_data['right_trigger_value'] = right_controller_data[3]
                    controller_data['right_touchpad'] = bool(right_controller_data[4] > 0.5)
                    controller_data['right_touchpad_x'] = right_controller_data[5]
                    controller_data['right_touchpad_y'] = right_controller_data[6]
                    controller_data['right_thumbstick'] = bool(right_controller_data[7] > 0.5)
                    controller_data['right_thumbstick_x'] = right_controller_data[8]
                    controller_data['right_thumbstick_y'] = right_controller_data[9]
                    controller_data['right_a_button'] = bool(right_controller_data[10] > 0.5)
                    controller_data['right_b_button'] = bool(right_controller_data[11] > 0.5)
                    
                    # 格式化打印右手柄状态
                    right_status = (
                        f"------111 右手柄状态:\n"
                        f"  侧握键: {controller_data['right_squeeze']} (值: {controller_data['right_squeeze_value']:.2f})\n"
                        f"  扣机键: {controller_data['right_trigger']} (值: {controller_data['right_trigger_value']:.2f})\n"
                        f"  触控板: {controller_data['right_touchpad']} (X: {controller_data['right_touchpad_x']:.2f}, Y: {controller_data['right_touchpad_y']:.2f})\n"
                        f"  摇杆: {controller_data['right_thumbstick']} (X: {controller_data['right_thumbstick_x']:.2f}, Y: {controller_data['right_thumbstick_y']:.2f})\n"
                        f"  A按钮: {controller_data['right_a_button']}\n"
                        f"  B按钮: {controller_data['right_b_button']}"
                    )
                    logger.info(right_status)
                else:
                    # 数据长度不足，记录警告
                    logger.warning(f"------111 右手柄数据长度不足: {len(right_controller_data)}")
                    # 保持默认值
                
                # 检查左右手腕标志一致性
                if left_wrist_flag != right_wrist_flag and left_wrist_flag == False:
                    # logger.warning("------111 左右手腕标志不一致，右手柄侧握键未被按下")
                    # logger.warning(f"------111 left_wrist_flag: {left_wrist_flag}, right_wrist_flag: {right_wrist_flag}")
                    controller_data['right_a_button'] = False
                   
                else:
                    logger.info(f"------111 左右手腕标志一致，右手柄侧握键, left_wrist: {left_wrist_flag}, right_wrist: {right_wrist_flag}, right_squeeze: {controller_data['right_squeeze']}")
            
            # 获取左手柄状态
            if hasattr(self.tv, 'left_controller'):
                left_controller_data = self.tv.left_controller()
                # 获取完整的左手柄状态
                if len(left_controller_data) >= 12:  # 确保数组足够长
                    # 基本按钮
                    controller_data['left_squeeze'] = bool(left_controller_data[0] > 0.5)
                    controller_data['left_squeeze_value'] = left_controller_data[1]
                    # 其他按钮
                    controller_data['left_trigger'] = bool(left_controller_data[2] > 0.5)
                    controller_data['left_trigger_value'] = left_controller_data[3]
                    controller_data['left_touchpad'] = bool(left_controller_data[4] > 0.5)
                    controller_data['left_touchpad_x'] = left_controller_data[5]
                    controller_data['left_touchpad_y'] = left_controller_data[6]
                    controller_data['left_thumbstick'] = bool(left_controller_data[7] > 0.5)
                    controller_data['left_thumbstick_x'] = left_controller_data[8]
                    controller_data['left_thumbstick_y'] = left_controller_data[9]
                    controller_data['left_x_button'] = bool(left_controller_data[10] > 0.5)
                    controller_data['left_y_button'] = bool(left_controller_data[11] > 0.5)
                    
                    # 调试信息：打印原始数据
                    logger.debug(f"左手柄原始数据: {left_controller_data[:]}")
                    
                    # 格式化打印左手柄状态
                    left_status = (
                        f"------111 左手柄状态:\n"
                        f"  侧握键: {controller_data['left_squeeze']} (值: {controller_data['left_squeeze_value']:.2f})\n"
                        f"  扣机键: {controller_data['left_trigger']} (值: {controller_data['left_trigger_value']:.2f})\n"
                        f"  触控板: {controller_data['left_touchpad']} (X: {controller_data['left_touchpad_x']:.2f}, Y: {controller_data['left_touchpad_y']:.2f})\n"
                        f"  摇杆: {controller_data['left_thumbstick']} (X: {controller_data['left_thumbstick_x']:.2f}, Y: {controller_data['left_thumbstick_y']:.2f})\n"
                        f"  A按钮: {controller_data['left_x_button']}\n"
                        f"  B按钮: {controller_data['left_y_button']}\n"
                        f"  原始数据: A={left_controller_data[10]:.2f}, B={left_controller_data[11]:.2f}"
                    )
                    logger.info(left_status)
                else:
                    # 数据长度不足，记录警告
                    logger.warning(f"------111 左手柄数据长度不足: {len(left_controller_data)}")
                    # 保持默认值
        except Exception as e:
            logger.error(f"从共享内存获取控制器数据时出错: {e}")
        
        # 改变坐标基约定：将VuerMat（（坐标基）OpenXR约定）转换为WristMat（（坐标基）机器人约定）
        # 附注：WristMat = T_{robot}_{OpenXR} * VuerMat * T_{robot}_{OpenXR}^T
        # 右乘fast_mat_inv(T_robot_OpenXR)的原因：
        #   这是相似变换：B = PAP^{-1}，即B ~ A
        #   例如：
        #   - 对于机器人约定下的位姿数据T_r，左乘WristMat意味着：
        #   - WristMat * T_r  ==>  T_{robot}_{OpenXR} * VuerMat * T_{OpenXR}_{robot} * T_r
        #   - 首先，转换到OpenXR约定（T_{OpenXR}_{robot}的作用）
        #   - 然后，在OpenXR约定下应用旋转VuerMat（VuerMat的作用）
        #   - 最后，转换回机器人约定（T_{robot}_{OpenXR}的作用）
        #   这使得在机器人约定下的旋转效果与在OpenXR约定下相同。实现将VR坐标系是xyz:右上后，转换为机器人坐标系xyz:前左上。
        head_mat = T_robot_openxr @ head_vuer_mat @ fast_mat_inv(T_robot_openxr)
        left_wrist_mat  = T_robot_openxr @ left_wrist_vuer_mat @ fast_mat_inv(T_robot_openxr)
        right_wrist_mat = T_robot_openxr @ right_wrist_vuer_mat @ fast_mat_inv(T_robot_openxr)

        logger.info(f"------222 robot head_mat:\n{print_pose_info(head_mat)}")
        logger.info(f"------222 robot left_wrist_mat:\n{print_pose_info(left_wrist_mat)}")
        logger.info(f"------222 robot right_wrist_mat:\n{print_pose_info(right_wrist_mat)}")

        # 转换为机器人坐标系xyz:前左上后，
        # 绕自身Y轴逆时针旋转95度，然后在运动坐标系下绕Z轴逆时针旋转180度，得到与手腕末端坐标一致的坐标。
        unitree_left_wrist = left_wrist_mat @ T_to_fangnuo_left_wrist
        unitree_right_wrist = right_wrist_mat @ T_to_fangnuo_right_wrist
 
     
        # 从世界坐标转换为头部坐标（仅平移）。
        unitree_left_wrist[0:3, 3]  = unitree_left_wrist[0:3, 3] - head_mat[0:3, 3]
        unitree_right_wrist[0:3, 3] = unitree_right_wrist[0:3, 3] - head_mat[0:3, 3]
       
        logger.info(f"-----333 robot head unitree_left_wrist:\n{print_pose_info(unitree_left_wrist)}")
        logger.info(f"-----333 robot head unitree_right_wrist:\n{print_pose_info(unitree_right_wrist)}")
        

        # --------------------------------hand-------------------------------------

        # 齐次坐标，将[xyz]转换为[xyz1]
        # 附注：np.concatenate([25,3]^T,(1,25)) ==> hand_vuer_mat的形状为(4,25)
        # 现在根据（基本的）OpenXR规范，矩阵形状如下：
        #    x0 x1 x2 ··· x23 x24
        #    y0 y1 y1 ··· y23 y24
        #    z0 z1 z2 ··· z23 z24
        #     1  1  1 ···   1   1
        # 获取手部坐标数据，先调用方法获取数据，然后再使用copy()
        left_landmarks_data = self.tv.left_landmarks()
        right_landmarks_data = self.tv.right_landmarks()
        
        left_hand_vuer_mat  = np.concatenate([left_landmarks_data.copy().T, np.ones((1, left_landmarks_data.shape[0]))])
        right_hand_vuer_mat = np.concatenate([right_landmarks_data.copy().T, np.ones((1, right_landmarks_data.shape[0]))])

        
        # 改变坐标基约定：从（基）OpenXR 约定到（基）机器人约定
        # 只是三维点的坐标基改变。无旋转，只有平移。无需右乘 fast_mat_inv(T_robot_openxr)。
        left_hand_mat  = T_robot_openxr @ left_hand_vuer_mat
        right_hand_mat = T_robot_openxr @ right_hand_vuer_mat


        # 从世界坐标系转换到手腕坐标系。（此过程遵循（基本）机器人惯例）
        # 附注：基于手腕的手部矩阵 HandMat_WristBased = 手腕矩阵 WristMat_{世界}_{手腕}^T * 手部矩阵 HandMat_{世界}
        #       基于手腕的手部矩阵 HandMat_WristBased = 手腕矩阵 WristMat_{手腕}_{世界}   * 手部矩阵 HandMat_{世界}，即手部矩阵 HandMat_{手腕}
        left_hand_mat_wb  = fast_mat_inv(left_wrist_mat) @ left_hand_mat
        right_hand_mat_wb = fast_mat_inv(right_wrist_mat) @ right_hand_mat

        # 更改手部约定：将HandMat（（左手）XR/Apple Vision Pro约定）转换为UnitreeHandMat（（左手URDF）宇树约定）
        # 左乘的原因：T_to_fangnuo_hand @ left_hand_mat_wb ==> （4,4） @ （4,25） ==> （4,25），（4,25）[0:3, :] ==> （3,25），（3,25）.T ==> （25,3）           
        # 现在在（左手URDF）宇树约定下，矩阵形状如下：
        #    [x0, y0, z0]
        #    [x1, y1, z1]
        #    ···
        #    [x23,y23,z23] 
        #    [x24,y24,z24]               
        unitree_left_hand  = (T_to_fangnuo_hand @ left_hand_mat_wb)[0:3, :].T
        unitree_right_hand = (T_to_fangnuo_hand @ right_hand_mat_wb)[0:3, :].T

        # --------------------------------offset-------------------------------------

        head_rmat = head_mat[:3, :3]
        
        # 逆运动学求解的坐标原点是腰部关节电机。你可以使用teleop/robot_control/robot_arm_ik.py单元测试来检查这一点。
        # unitree左腕部的坐标原点是头部。因此，有必要将unitree左腕部的坐标原点从头部转换到腰部。
        # scale_arms 缩放了动作，保证初始化腕部位姿可到达
        unitree_left_wrist[0, 3] +=0.2 # x=0.13
        unitree_left_wrist[1, 3] +=0.15 # y=0.25   
        unitree_left_wrist[2, 3] +=(0.5 + 0.78) # z=0.78
        unitree_right_wrist[0,3] +=0.2
        unitree_right_wrist[1,3] -=0.15
        unitree_right_wrist[2,3] +=(0.5 + 0.78)
        logger.info(f"------444 final left_wrist:\n{print_pose_info(unitree_left_wrist)}")
        logger.info(f"------444 final right_wrist:\n{print_pose_info(unitree_right_wrist)}")

        return head_rmat, unitree_left_wrist, unitree_right_wrist, unitree_left_hand, unitree_right_hand, controller_data